import cv2
import numpy as np

from yolop_model.yolop import get_net
from loaddata import VideoReader
import torch
from tqdm import tqdm
import random
import torchvision.transforms as transforms
normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
    )

transform = transforms.Compose([
            transforms.ToTensor(),
            normalize,
        ])
def pred():
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model = get_net()
    checkpoint = torch.load('weights/yolop.pth', map_location= device)
    model.load_state_dict(checkpoint['state_dict'])
    model = model.to(device)
    model.eval()
    half = device.type != 'cpu'
    if half:
        model.half()  # to FP16
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]
    videodataset = VideoReader(path='cope_data/')
    # print(model)
    with torch.no_grad():
        for i,(path, img, img_det, vid_cap, shapes) in tqdm(enumerate(videodataset),total=len(videodataset)):
            img = transform(img).to(device)
            img = img.half() if half else img.float()  # uint8 to fp16/32
            if img.ndimension() == 3:
                img = img.unsqueeze(0)
            # Inference
            det_out, da_seg_out,ll_seg_out = model(img)
            # print(len(det_out),da_seg_out.size(),ll_seg_out.size())


            _, _, height, width = img.shape
            h, w, _ = img_det.shape
            pad_w, pad_h = shapes[1][1]
            pad_w = int(pad_w)
            pad_h = int(pad_h)
            ratio = shapes[1][0][1]

            da_predict = da_seg_out[:, :, pad_h:(height - pad_h), pad_w:(width - pad_w)]
            # print(da_predict.size())
            da_seg_mask = torch.nn.functional.interpolate(da_predict, scale_factor=int(1 / ratio), mode='bilinear')
            _, da_seg_mask = torch.max(da_seg_mask, 1)
            da_seg_mask = da_seg_mask.int().squeeze().cpu().numpy()
            # da_seg_mask = morphological_process(da_seg_mask, kernel_size=7)
            # print(da_seg_mask.shape)

            ll_predict = ll_seg_out[:, :, pad_h:(height - pad_h), pad_w:(width - pad_w)]
            # print(ll_predict.size())
            ll_seg_mask = torch.nn.functional.interpolate(ll_predict, scale_factor=int(1 / ratio), mode='bilinear')
            _, ll_seg_mask = torch.max(ll_seg_mask, 1)
            ll_seg_mask = ll_seg_mask.int().squeeze().cpu().numpy()
            # print(ll_seg_mask.shape)


            da_seg_mask = np.array(da_seg_mask*255, dtype=np.uint8)
            ll_seg_mask = np.array(ll_seg_mask*255, dtype=np.uint8)
            pic = np.vstack((da_seg_mask[::2,::2],ll_seg_mask[::2,::2]))
            cv2.imshow('pic',pic)
            cv2.waitKey(0)

            # da_seg_mask = cv2.resize(da_seg_mask,(1280,720))
            # print(da_seg_mask.shape, da_seg_mask[3 * 720 // 4 + 20, 4 * 1280 // 9 + 10])
            # ret = np.where(np.logical_and(da_seg_mask>0,da_seg_mask<255))
            # print(ret)
            # ret = np.where(np.logical_and(ll_seg_mask>0,ll_seg_mask<255))
            # print(ret)

            # break
        cv2.destroyAllWindows()
        print(names)
        print(colors)
if __name__ == '__main__':
    pred()